'''
Soft Actor-Critic version 2
using target Q instead of V net: 2 Q net, 2 target Q net, 1 policy net
add alpha loss compared with version 1
paper: https://arxiv.org/pdf/1812.05905.pdf

Discrete version reference: 
https://towardsdatascience.com/adapting-soft-actor-critic-for-discrete-action-spaces-a20614d4a50a

sac algo code reference:
https://github.com/quantumiracle/Popular-RL-Algorithms.git


reward 需要调整大小接近 更新的后半部分 /2000.
'''

import random
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical
from IPython.display import clear_output
import matplotlib.pyplot as plt
import argparse
import wandb

from env.env2_v2 import MECEnv

GPU = True
device_idx = 0
if GPU:
    device = torch.device("cuda:" + str(device_idx) if torch.cuda.is_available() else "cpu")
else:
    device = torch.device("cpu")
print(device)

parser = argparse.ArgumentParser(description='Train or test neural net motor controller.')
parser.add_argument('--train', dest='train', action='store_true', default=True)
parser.add_argument('--test', dest='test', action='store_true', default=False)

args = parser.parse_args()

user_num = 4
server_num = 4
# LR = 6e-4 #6e-4 # 4e-5
LR = 5e-3


def norm_state(state):
    return state / torch.tensor([[[3000, 4, 1, 5, 1, 20e5, 2]]]).to(device)
    # return state


def reward_scaler(reward, reward_scale=1):
    # std = reward.std(dim=0)
    # mean = reward.mean(dim=0)
    # reward = reward_scale * (reward - mean) / (std)

    reward = reward * 10 + 0.9
    # reward = reward + 0.85
    # reward = reward / 2.0 + 0.3
    # reward = reward / 18.0 * 10 + 0.3
    return reward


class ReplayBuffer:
    def __init__(self, capacity):
        self.capacity = capacity
        self.buffer = []
        self.position = 0

    def push(self, state, action, reward, next_state, done):
        if len(self.buffer) < self.capacity:
            self.buffer.append(None)
        self.buffer[self.position] = (state, action, reward, next_state, done)
        self.position = int((self.position + 1) % self.capacity)  # as a ring buffer

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        state, action, reward, next_state, done = map(np.stack, zip(*batch))  # stack for each element
        ''' 
        the * serves as unpack: sum(a,b) <=> batch=(a,b), sum(*batch) ;
        zip: a=[1,2], b=[2,3], zip(a,b) => [(1, 2), (2, 3)] ;
        the map serves as mapping the function on each list element: map(square, [2,3]) => [4,9] ;
        np.stack((1,2)) => array([1, 2])
        '''
        return state, action, reward, next_state, done

    def __len__(self):
        return len(self.buffer)


class SoftQNetwork(nn.Module):
    def __init__(self, num_inputs, num_actions, hidden_size, init_w=3e-3):
        super(SoftQNetwork, self).__init__()

        self.linear1 = nn.Linear(num_inputs, hidden_size)
        # self.linear2 = nn.Linear(hidden_size, hidden_size)
        self.linear3 = nn.Linear(hidden_size, hidden_size)
        self.linear4 = nn.Linear(hidden_size, num_actions)

        self.linear4.weight.data.uniform_(-init_w, init_w)
        self.linear4.bias.data.uniform_(-init_w, init_w)

    def forward(self, state):
        state = norm_state(state)
        x = F.relu(self.linear1(state))
        # x = F.relu(self.linear2(x))
        x = F.tanh(self.linear3(x))
        # x = F.tanh(self.linear3(x))
        x = self.linear4(x)
        return x


class PolicyNetwork(nn.Module):
    def __init__(self, num_inputs, num_actions, hidden_size, init_w=3e-3, log_std_min=-20, log_std_max=2):
        super(PolicyNetwork, self).__init__()

        self.linear1 = nn.Linear(num_inputs, hidden_size)
        # self.linear2 = nn.Linear(hidden_size, hidden_size)
        # self.linear3 = nn.Linear(hidden_size, hidden_size)
        self.linear4 = nn.Linear(hidden_size, hidden_size)

        self.output = nn.Linear(hidden_size, num_actions)

        self.num_actions = num_actions

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.xavier_uniform_(module.weight)

    def forward(self, state, softmax_dim=-1):
        state = norm_state(state)
        x = F.relu(self.linear1(state))
        x = F.tanh(self.linear4(x))
        out = self.output(x)

        probs = F.softmax(out, dim=softmax_dim)
        return probs

    def evaluate(self, state, epsilon=1e-8):
        '''
        generate sampled action with state as input wrt the policy network;
        '''
        probs = self.forward(state, softmax_dim=-1)
        # print(probs)
        log_probs = torch.log(probs)

        # Avoid numerical instability. Ref: https://github.com/ku2482/sac-discrete.pytorch/blob/40c9d246621e658750e0a03001325006da57f2d4/sacd/model.py#L98
        z = (probs == 0.0).float() * epsilon
        log_probs = torch.log(probs + z)

        return log_probs

    def get_action(self, state, deterministic):
        state = torch.FloatTensor(state).unsqueeze(0).to(device)
        probs = self.forward(state)
        dist = Categorical(probs)  # 根据概率采样

        if deterministic:
            detach = probs.detach()
            nump = detach.cpu().numpy()
            action = np.argmax(nump, axis=2)[0]
        else:
            action = dist.sample().cpu().numpy()[0]
            # action = action.squeeze()
            # action = action.detach().cpu()
            # action = action.view(-1)
        return action


class SAC_Trainer():
    def __init__(self, replay_buffer, hidden_dim):
        self.replay_buffer = replay_buffer

        self.soft_q_net1 = SoftQNetwork(state_dim, action_dim, hidden_dim).to(device)
        self.soft_q_net2 = SoftQNetwork(state_dim, action_dim, hidden_dim).to(device)
        self.target_soft_q_net1 = SoftQNetwork(state_dim, action_dim, hidden_dim).to(device)
        self.target_soft_q_net2 = SoftQNetwork(state_dim, action_dim, hidden_dim).to(device)
        self.policy_net = PolicyNetwork(state_dim, action_dim, hidden_dim).to(device)
        self.log_alpha = torch.zeros(1, dtype=torch.float32, requires_grad=True, device=device)
        print('Soft Q Network (1,2): ', self.soft_q_net1)
        print('Policy Network: ', self.policy_net)

        for target_param, param in zip(self.target_soft_q_net1.parameters(), self.soft_q_net1.parameters()):
            target_param.data.copy_(param.data)
        for target_param, param in zip(self.target_soft_q_net2.parameters(), self.soft_q_net2.parameters()):
            target_param.data.copy_(param.data)

        self.soft_q_criterion1 = nn.MSELoss()
        self.soft_q_criterion2 = nn.MSELoss()

        # 1,6 1e-4 在1k左右收敛到最高
        # 6,6 5e-4

        policy_lr = LR  # 3e-3  10 ,5   # w0的时候5e-4 收敛到 -20
        soft_q_lr = LR
        alpha_lr = LR

        self.soft_q_optimizer1 = optim.Adam(self.soft_q_net1.parameters(), lr=soft_q_lr)
        self.soft_q_optimizer2 = optim.Adam(self.soft_q_net2.parameters(), lr=soft_q_lr)
        self.policy_optimizer = optim.Adam(self.policy_net.parameters(), lr=policy_lr)
        self.alpha_optimizer = optim.Adam([self.log_alpha], lr=alpha_lr)

    def update(self, step, batch_size, reward_scale=10., auto_entropy=True, target_entropy=-2, gamma=0.99,
               soft_tau=1e-2, ):
        state, action, reward, next_state, done = self.replay_buffer.sample(batch_size)
        # print('sample:', state, action,  reward, done)

        state = torch.FloatTensor(state).to(device)
        next_state = torch.FloatTensor(next_state).to(device)
        action = torch.Tensor(action).to(torch.int64).to(device)
        reward = torch.FloatTensor(reward).unsqueeze(1).unsqueeze(-1).to(
            device)  # reward is single value, unsqueeze() to add one dim to be [reward] at the sample dim;
        done = torch.FloatTensor(np.float32(done)).unsqueeze(1).unsqueeze(-1).to(device)
        predicted_q_value1 = self.soft_q_net1(state)
        predicted_q_value1 = predicted_q_value1.gather(2, action.unsqueeze(-1))
        predicted_q_value2 = self.soft_q_net2(state)
        predicted_q_value2 = predicted_q_value2.gather(2, action.unsqueeze(-1))
        log_prob = self.policy_net.evaluate(state)
        with torch.no_grad():
            next_log_prob = self.policy_net.evaluate(next_state)
        # reward = reward_scale * (reward - reward.mean(dim=0)) / (reward.std(dim=0) + 1e-6) # normalize with batch mean and std; plus a small number to prevent numerical problem
        reward = reward_scaler(reward)

        # Training Q Function
        self.alpha = self.log_alpha.exp()
        # target_q_min = (next_log_prob.exp() * (torch.min(self.target_soft_q_net1(next_state), self.target_soft_q_net2(next_state)) - self.alpha * next_log_prob)).sum(dim=-1).unsqueeze(-1)
        _ = (next_log_prob.exp() * (torch.min(self.target_soft_q_net1(next_state),
                                              self.target_soft_q_net2(next_state)) - self.alpha * next_log_prob))
        target_q_min = _.sum(dim=-1).unsqueeze(-1)
        # reward是不是太高了  设置调整reward的影响
        target_q_value = reward + (1 - done) * gamma * target_q_min  # if done==1, only reward
        q_value_loss1 = self.soft_q_criterion1(predicted_q_value1,
                                               target_q_value.detach())  # detach: no gradients for the variable
        q_value_loss2 = self.soft_q_criterion2(predicted_q_value2, target_q_value.detach())

        self.soft_q_optimizer1.zero_grad()
        q_value_loss1.backward()
        self.soft_q_optimizer1.step()
        self.soft_q_optimizer2.zero_grad()
        q_value_loss2.backward()
        self.soft_q_optimizer2.step()

        # Training Policy Function
        with torch.no_grad():
            predicted_new_q_value = torch.min(self.soft_q_net1(state), self.soft_q_net2(state))
        policy_loss = (log_prob.exp() * (self.alpha * log_prob - predicted_new_q_value)).sum(dim=-1).mean()
        self.policy_optimizer.zero_grad()
        policy_loss.backward()
        self.policy_optimizer.step()

        # Updating alpha wrt entropy
        # alpha = 0.0  # trade-off between exploration (max entropy) and exploitation (max Q) 
        if auto_entropy is True:
            alpha_loss = -(self.log_alpha * (log_prob + target_entropy).detach()).mean()
            # print('alpha loss: ',alpha_loss)
            self.alpha_optimizer.zero_grad()
            alpha_loss.backward()
            self.alpha_optimizer.step()
        else:
            self.alpha = 1.
            alpha_loss = 0

        # print('q loss: ', q_value_loss1.item(), q_value_loss2.item())
        # print('policy loss: ', policy_loss.item() )

        # Soft update the target value net
        for target_param, param in zip(self.target_soft_q_net1.parameters(), self.soft_q_net1.parameters()):
            target_param.data.copy_(  # copy data value into target parameters
                target_param.data * (1.0 - soft_tau) + param.data * soft_tau
            )
        for target_param, param in zip(self.target_soft_q_net2.parameters(), self.soft_q_net2.parameters()):
            target_param.data.copy_(  # copy data value into target parameters
                target_param.data * (1.0 - soft_tau) + param.data * soft_tau
            )

        return policy_loss.item(), q_value_loss1.item(), q_value_loss2.item(), predicted_new_q_value.mean().item()

    def save_model(self, path):
        torch.save(self.soft_q_net1.state_dict(), path + '_q1')
        torch.save(self.soft_q_net2.state_dict(), path + '_q2')
        torch.save(self.policy_net.state_dict(), path + '_policy')

    def load_model(self, path):
        self.soft_q_net1.load_state_dict(torch.load(path + '_q1'))
        self.soft_q_net2.load_state_dict(torch.load(path + '_q2'))
        self.policy_net.load_state_dict(torch.load(path + '_policy'))

        self.soft_q_net1.eval()
        self.soft_q_net2.eval()
        self.policy_net.eval()


def plot(rewards):
    clear_output(True)
    plt.figure(figsize=(20, 5))
    plt.plot(rewards)
    plt.savefig('sac_v2.png')
    plt.show()


replay_buffer_size = 10000
replay_buffer = ReplayBuffer(replay_buffer_size)

# choose env
env = MECEnv.env(user_num, server_num, 1.0, 0)

state_dim = env.state_dim
action_dim = env.action_dim  # discrete

# hyper-parameters for RL training
max_episodes = 300
max_steps = 2000
frame_idx = 0
batch_size = 128
update_itr = 1
AUTO_ENTROPY = True
DETERMINISTIC = False
hidden_dim = 128
rewards = []
model_path = './model/sac_discrete_v2'
target_entropy = -1 * action_dim
# target_entropy = 0.98 * -np.log(1 / action_dim)

sac_trainer = SAC_Trainer(replay_buffer, hidden_dim=hidden_dim)

if __name__ == '__main__':
    with wandb.init(project="PomdpEnv_1_0", name=f"sacdm_mdp_{user_num}_{server_num}_{LR}"):
        if args.train:
            # training loop
            for eps in range(max_episodes):
                state, _ = env.reset()
                episode_reward = 0
                ep_delay = 0
                ep_consumption = 0
                ep_policy_loss, ep_q_value_loss1, ep_q_value_loss2, ep_value = 0, 0, 0, 0
                for step in range(max_steps):
                    action = sac_trainer.policy_net.get_action(state, deterministic=DETERMINISTIC)
                    # self.state, self.reward, False, {}, delay, consumption, self.observation
                    next_state, reward, done, _, delay, consumption, _ = env.step(action)
                    # if eps % 20 == 0 and eps > 0:
                    #     print(next_state)
                    # env.render()

                    replay_buffer.push(state, action, reward, next_state, done)

                    state = next_state
                    episode_reward += reward
                    frame_idx += 1

                    # if len(replay_buffer) > batch_size:
                    if len(replay_buffer) > batch_size and step % 10 == 0:
                        for i in range(update_itr):
                            policy_loss, q_value_loss1, q_value_loss2, value = sac_trainer.update(step, batch_size,
                                                                                                  reward_scale=1.,
                                                                                                  auto_entropy=AUTO_ENTROPY,
                                                                                                  target_entropy=target_entropy,
                                                                                                  )
                            ep_policy_loss += policy_loss
                            ep_q_value_loss1 += q_value_loss1
                            ep_q_value_loss2 += q_value_loss2
                            ep_value += value

                    if done:
                        break

                # plot(rewards)
                # np.save('rewards', rewards)
                print('Episode: ', eps, '| Episode Reward: ', episode_reward, '| Episode Length: ', step,
                      '| policy_loss: ',
                      ep_policy_loss
                      , '| q_value_loss1: ', ep_q_value_loss1, '|q_value_loss2: ', ep_q_value_loss2, "|value:",
                      ep_value)

                wandb.log({"Episode Reward": episode_reward, "Policy Loss": ep_policy_loss,"delay": ep_delay,
                                   "energy consumption": ep_consumption })
                rewards.append(episode_reward)

            # sac_trainer.save_model(model_path)
